{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## 掩码多头注意力",
   "id": "a256d2d8a29856de"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-09T14:03:31.765158Z",
     "start_time": "2025-07-09T14:03:31.760634Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch\n",
    "\n",
    "# x表示一个输入的seq，decoder的输入\n",
    "x = torch.rand(4, 3)\n",
    "x"
   ],
   "id": "883e8e88f7982f8d",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.0544, 0.3815, 0.2969],\n",
       "        [0.2863, 0.1421, 0.0358],\n",
       "        [0.1022, 0.6532, 0.7013],\n",
       "        [0.1904, 0.3571, 0.0432]])"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 28
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-09T14:03:39.429628Z",
     "start_time": "2025-07-09T14:03:39.425707Z"
    }
   },
   "cell_type": "code",
   "source": [
    "wq = torch.rand(3, 3)\n",
    "wk = torch.rand(3, 3)\n",
    "wv = torch.rand(3, 3)"
   ],
   "id": "59939c6f166b34ca",
   "outputs": [],
   "execution_count": 29
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-09T14:03:43.187001Z",
     "start_time": "2025-07-09T14:03:43.182632Z"
    }
   },
   "cell_type": "code",
   "source": [
    "q = torch.matmul(x, wq)\n",
    "k = torch.matmul(x, wk)\n",
    "v = torch.matmul(x, wv)\n",
    "q, k, v"
   ],
   "id": "756b788566bd3f2f",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[0.5825, 0.1260, 0.5078],\n",
       "         [0.3939, 0.3009, 0.4188],\n",
       "         [1.1561, 0.2283, 0.9273],\n",
       "         [0.4851, 0.2471, 0.5458]]),\n",
       " tensor([[0.4276, 0.4159, 0.4140],\n",
       "         [0.3454, 0.3930, 0.1762],\n",
       "         [0.8540, 0.7932, 0.7470],\n",
       "         [0.3823, 0.4413, 0.3617]]),\n",
       " tensor([[0.3155, 0.4941, 0.3683],\n",
       "         [0.1278, 0.2936, 0.2272],\n",
       "         [0.6567, 0.9413, 0.7742],\n",
       "         [0.1844, 0.4283, 0.2515]]))"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 30
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-09T14:04:15.625281Z",
     "start_time": "2025-07-09T14:04:15.621971Z"
    }
   },
   "cell_type": "code",
   "source": [
    "qk = torch.matmul(q, k.T)\n",
    "qk"
   ],
   "id": "65399856fb7a9b9c",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.5117, 0.3401, 0.9767, 0.4619],\n",
       "        [0.4670, 0.3281, 0.8879, 0.4349],\n",
       "        [0.9733, 0.6524, 1.8611, 0.8781],\n",
       "        [0.5362, 0.3608, 1.0180, 0.4919]])"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 31
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-09T14:05:40.560889Z",
     "start_time": "2025-07-09T14:05:40.557948Z"
    }
   },
   "cell_type": "code",
   "source": "attention_score = qk / torch.sqrt(torch.tensor(3))",
   "id": "b6ea73fa8611b184",
   "outputs": [],
   "execution_count": 36
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-09T14:05:41.567367Z",
     "start_time": "2025-07-09T14:05:41.563983Z"
    }
   },
   "cell_type": "code",
   "source": "attention_score",
   "id": "6bb3d7276977371c",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.2954, 0.1964, 0.5639, 0.2667],\n",
       "        [0.2696, 0.1894, 0.5126, 0.2511],\n",
       "        [0.5619, 0.3767, 1.0745, 0.5070],\n",
       "        [0.3096, 0.2083, 0.5877, 0.2840]])"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 37
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-09T14:06:57.478274Z",
     "start_time": "2025-07-09T14:06:57.473655Z"
    }
   },
   "cell_type": "code",
   "source": [
    "mask = torch.tril(torch.ones(4, 4))\n",
    "mask"
   ],
   "id": "d31f0e84b66b90c4",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 0., 0., 0.],\n",
       "        [1., 1., 0., 0.],\n",
       "        [1., 1., 1., 0.],\n",
       "        [1., 1., 1., 1.]])"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 38
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-09T14:08:14.974455Z",
     "start_time": "2025-07-09T14:08:14.970805Z"
    }
   },
   "cell_type": "code",
   "source": [
    "attention_score = attention_score.masked_fill(mask == 0, float('-inf'))\n",
    "attention_score"
   ],
   "id": "6054dc9a1b4eca81",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.2954,   -inf,   -inf,   -inf],\n",
       "        [0.2696, 0.1894,   -inf,   -inf],\n",
       "        [0.5619, 0.3767, 1.0745,   -inf],\n",
       "        [0.3096, 0.2083, 0.5877, 0.2840]])"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 39
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-09T14:08:50.220420Z",
     "start_time": "2025-07-09T14:08:50.215819Z"
    }
   },
   "cell_type": "code",
   "source": [
    "attention_weight = torch.softmax(attention_score, dim=-1)\n",
    "attention_weight"
   ],
   "id": "6cb0eb65868a385e",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1.0000, 0.0000, 0.0000, 0.0000],\n",
       "        [0.5200, 0.4800, 0.0000, 0.0000],\n",
       "        [0.2857, 0.2374, 0.4770, 0.0000],\n",
       "        [0.2381, 0.2152, 0.3145, 0.2321]])"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 40
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-09T14:09:08.147447Z",
     "start_time": "2025-07-09T14:09:08.143581Z"
    }
   },
   "cell_type": "code",
   "source": [
    "o = torch.matmul(attention_weight, v)  # qk的形状是(4, 4), v的形状是(4, 2), o的形状是(4, 2), 输出是q的行，v的列, q的行等于i的行，这就对上了\n",
    "o"
   ],
   "id": "6c950b653bd20aea",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.3155, 0.4941, 0.3683],\n",
       "        [0.2254, 0.3978, 0.3006],\n",
       "        [0.4337, 0.6598, 0.5284],\n",
       "        [0.3520, 0.5763, 0.4385]])"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 41
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-09T14:24:16.732438Z",
     "start_time": "2025-07-09T14:24:16.727045Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "class MultiHeadAttention(nn.Module):\n",
    "\n",
    "    def __init__(self, embed_dim: int, attn_dim: int, output_dim: int, num_heads: int):\n",
    "        super().__init__()\n",
    "\n",
    "        self.embed_dim = embed_dim\n",
    "        self.attn_dim = attn_dim\n",
    "        self.output_dim = output_dim\n",
    "        self.num_heads = num_heads\n",
    "        self.head_dim = attn_dim // num_heads # //表示向下取整，attn_dim是head_dim的整数倍\n",
    "\n",
    "        # QKV投影层：从输入维度映射到内部维度\n",
    "        # projection\n",
    "        self.q_proj = nn.Linear(embed_dim, self.attn_dim)\n",
    "        self.k_proj = nn.Linear(embed_dim, self.attn_dim)\n",
    "        self.v_proj = nn.Linear(embed_dim, self.attn_dim)\n",
    "\n",
    "        # 输出投影层：从内部维度映射到输出维度\n",
    "        self.out_proj = nn.Linear(self.attn_dim, self.output_dim)\n",
    "\n",
    "    def forward(self, x, mask=None):\n",
    "        \"\"\"\n",
    "        输入: [batch_size, seq_len, embed_dim]\n",
    "        返回: [batch_size, seq_len, output_dim]\n",
    "        \"\"\"\n",
    "        batch_size, seq_len, embed_dim = x.shape\n",
    "\n",
    "        # 投影到QKV空间\n",
    "        q = self.q_proj(x)  # [batch_size, seq_len, attn_dim]\n",
    "        k = self.k_proj(x)  # [batch_size, seq_len, attn_dim]\n",
    "        v = self.v_proj(x)  # [batch_size, seq_len, attn_dim]\n",
    "\n",
    "        # [batch_size, seq_len, num_heads, head_dim]\n",
    "        # 分割多头 [batch_size, num_heads, seq_len, head_dim]\n",
    "        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)\n",
    "        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)\n",
    "        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)\n",
    "\n",
    "        # 计算注意力得分\n",
    "        # q   [batch_size, num_heads, seq_len, head_dim]\n",
    "        # k.T [batch_size, num_heads, head_dim, seq_len]\n",
    "        # q @ k.T 形状: [batch_size, num_heads, seq_len, seq_len]\n",
    "        attn_scores = torch.matmul(q, k.transpose(-2, -1))\n",
    "\n",
    "        # 缩放因子：防止乘积过大\n",
    "        d_k = k.size(-1)\n",
    "        attn_scores = attn_scores / torch.sqrt(torch.tensor(d_k))\n",
    "\n",
    "        if mask is not None:\n",
    "            attn_scores = attn_scores.masked_fill(mask == 0, float('-inf'))\n",
    "\n",
    "        # 计算注意力权重\n",
    "        attn_weights = torch.softmax(attn_scores, dim=-1)\n",
    "\n",
    "        # 计算注意力输出\n",
    "        # attn_weights [batch_size, num_heads, seq_len, seq_len]\n",
    "        # v            [batch_size, num_heads, seq_len, head_dim]\n",
    "        # [batch_size, num_heads, seq_len, head_dim]\n",
    "        attn_out = torch.matmul(attn_weights, v)\n",
    "\n",
    "        # 合并多头 [batch_size, seq_len, attn_dim]\n",
    "\n",
    "        # [batch_size, seq_len, num_heads, head_dim]\n",
    "        attn_out = attn_out.transpose(1, 2).reshape(batch_size, seq_len, self.attn_dim)\n",
    "\n",
    "        # 投影到输出空间\n",
    "        return self.out_proj(attn_out), attn_scores, attn_weights"
   ],
   "id": "23b49745311ff540",
   "outputs": [],
   "execution_count": 45
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-09T14:24:21.295554Z",
     "start_time": "2025-07-09T14:24:21.291126Z"
    }
   },
   "cell_type": "code",
   "source": [
    "x = torch.rand(1, 4, 2)\n",
    "\n",
    "mask = torch.tril(torch.ones(4, 4))\n",
    "\n",
    "attn = MultiHeadAttention(embed_dim=2, attn_dim=6, output_dim=8, num_heads=2)\n",
    "out, attn_scores, attn_weights = attn(x, mask=mask)\n",
    "\n",
    "print(x.shape)\n",
    "print(out.shape)"
   ],
   "id": "b1f07af4d09cf745",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 4, 2])\n",
      "torch.Size([1, 4, 8])\n"
     ]
    }
   ],
   "execution_count": 46
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-09T14:24:25.693862Z",
     "start_time": "2025-07-09T14:24:25.690544Z"
    }
   },
   "cell_type": "code",
   "source": "attn_scores",
   "id": "56529e75b65ec140",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[-0.5099,    -inf,    -inf,    -inf],\n",
       "          [-0.5785, -0.4869,    -inf,    -inf],\n",
       "          [-0.4910, -0.4737, -0.3818,    -inf],\n",
       "          [-0.6102, -0.4661, -0.4205, -0.6389]],\n",
       "\n",
       "         [[-0.1833,    -inf,    -inf,    -inf],\n",
       "          [-0.2779, -0.2538,    -inf,    -inf],\n",
       "          [-0.1444, -0.1680, -0.1133,    -inf],\n",
       "          [-0.3349, -0.3442, -0.2606, -0.4234]]]],\n",
       "       grad_fn=<MaskedFillBackward0>)"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 47
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-09T14:24:38.007375Z",
     "start_time": "2025-07-09T14:24:38.004092Z"
    }
   },
   "cell_type": "code",
   "source": "attn_weights",
   "id": "6b6fd5f6606c282a",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[1.0000, 0.0000, 0.0000, 0.0000],\n",
       "          [0.4771, 0.5229, 0.0000, 0.0000],\n",
       "          [0.3192, 0.3248, 0.3560, 0.0000],\n",
       "          [0.2307, 0.2664, 0.2788, 0.2241]],\n",
       "\n",
       "         [[1.0000, 0.0000, 0.0000, 0.0000],\n",
       "          [0.4940, 0.5060, 0.0000, 0.0000],\n",
       "          [0.3324, 0.3247, 0.3429, 0.0000],\n",
       "          [0.2511, 0.2487, 0.2704, 0.2298]]]], grad_fn=<SoftmaxBackward0>)"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 48
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
